
import os
import random
import struct

class MnistData(object):
	"""docstring for MnistData"""
	def __init__(self, img_file, lable_file):
		super(MnistData, self).__init__()
		self.img_file = img_file
		self.lable_file = lable_file
		self.img_row = None
		self.img_col = None
		self.lable_size = 1

		self.img_header_size = 16
		self.lable_header_size = 8

		self.img_fs = None
		self.lable_fs = None

		self.count = None
		self._open_file()


	def _open_file(self):
		if not os.path.exists(self.lable_file):
			script_abs_dir = os.path.dirname(__file__)
			script_abs_dir = os.path.abspath(script_abs_dir)
			if os.path.exists(os.path.join(script_abs_dir, self.lable_file)):
				self.lable_file = os.path.join(script_abs_dir, self.lable_file)
				self.img_file = os.path.join(script_abs_dir, self.img_file)

		self.lable_fs = open(self.lable_file, "rb")
		buf = self.lable_fs.read(self.lable_header_size)
		magic, count = struct.unpack('>ii', buf)
		if magic != 2049:
			raise Exception("read lable head error: %d"%(magic))
		self.img_fs = open(self.img_file, "rb")
		buf = self.img_fs.read(self.img_header_size)
		i_magic, i_count, row, col = struct.unpack('>iiii', buf)
		if i_magic != 2051:
			raise Exception("read img head error: %d"%(i_magic))
		if i_count != count:
			raise Exception("lable count[%d] not equal img count[%d]"%(count, i_count))
		self.count = count
		self.img_row = row
		self.img_col = col

	def read_data(self, fs, start_count, count, data_size, header_size):
		offset = data_size * start_count + header_size
		# print "pos: size: ", fs.tell()
		# print "offset: ", offset
		fs.seek(offset, os.SEEK_SET)
		# print "cur pos: size: ", fs.tell()
		read_count = count * data_size
		# rss = []
		# while read_count > 0:
		rs = fs.read(read_count)
		if len(rs) < read_count:
			raise Exception("read file error")
			# read_count -= len(rs)
			# rss.append(rs)
		# return ''.join(rss)
		return rs

	def _read_img(self, fs, img_inx, count, row, col, header_size):
		return self.read_data(fs, img_inx, count, row*col, header_size)

	def _read_lable(self, fs, inx, count, lable_size, header_size):
		return self.read_data(fs, inx, count, lable_size, header_size)

	def read_img(self, img_inx, count):
		rs = self._read_img(self.img_fs, img_inx, count, 
			self.img_row, self.img_col, self.img_header_size)
		imgmats = []
		col_f = 'B'*(self.img_col*self.img_row)
		# col_f = 'B'*self.img_col
		col_inx = 0
		for i in xrange(count):
			last_col_inx = col_inx + self.img_col*self.img_row
			img = struct.unpack(col_f, rs[col_inx: last_col_inx])
			img = [ round(pix/255.0, 10) for pix in img]
			col_inx = last_col_inx
			# img = []
			# for r in xrange(self.img_row):
			# 	last_col_inx = col_inx + self.img_col
			# 	td = struct.unpack(col_f, rs[col_inx: last_col_inx])

			# 	td = [ [ pix / 255.0 ] for pix in td]
			# 	col_inx = last_col_inx
			# 	img.append(td)
			imgmats.append(img)
		return imgmats

	def read_lable(self, inx, count):
		rs = self._read_lable(self.lable_fs, inx, count,
			self.lable_size, self.lable_header_size)
		lables = struct.unpack('B'*count, rs)
		map_lable = [[0 for _ in xrange(10)] for l in lables]
		for i in xrange(len(lables)):
			map_lable[i][lables[i]] = 1
		return map_lable

	def get_inx(self, n, max_inx):
		inx = int(random.random() * max_inx) - n -1
		if inx < 0: 
			return 0
		return inx

	def next(self, n):
		inx = self.get_inx(n, self.count)
		lables = self.read_lable(inx, n)
		imgs = self.read_img(inx, n)
		return imgs, lables

train = MnistData("data/train-images-idx3-ubyte", "data/train-labels-idx1-ubyte")
test = MnistData("data/t10k-images-idx3-ubyte", "data/t10k-labels-idx1-ubyte")

def main():
	print train.next(1)
	# print test.next(1)

if __name__ == '__main__':
	main()